import sys
import os
import pandas as pd

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import dragonquant as dq
from dragonquant.api import order, log
from dragonquant.indicators import calculate_sma, calculate_rsi, calculate_atr
from dragonquant.cerebro import Cerebro

def get_hs300_components_mock():
    """Simulates fetching a dynamic stock universe."""
    return [
        "000651.SZ",
        "600000.SH",
        "600036.SH",
        "600519.SH",
    ]

class SystematicRiskParityStrategy(dq.Strategy):
    """
    A systematic, risk-parity mid-term strategy.
    - Trades a stock universe based on a trend-following and pullback entry system.
    - Manages risk by allocating equal risk exposure to each trade via ATR-based position sizing.
    """
    params = {
        'sma_long_window': 150, 'rsi_window': 14, 'rsi_pullback_level': 50,
        'rsi_profit_target': 70, 'atr_window': 20, 'atr_stop_multiplier': 2.5,
        'risk_per_trade_percent': 0.015, 'max_active_positions': 5,
    }

    def init(self):
        self.positions_data = {}
        if not self.universe: raise ValueError("Strategy requires a non-empty universe.")
        log(f"--- Strategy '{self.__class__.__name__}' Initialized ---")
        log(f"Risk per trade: {self.params['risk_per_trade_percent']:.2%}, Max positions: {self.params['max_active_positions']}")

    def next(self, context: Cerebro):
        for stock in list(self.positions_data.keys()):
            self._handle_exit_for_stock(stock, context)

        if len(self.positions_data) >= self.params['max_active_positions']: return
        
        for stock in context.universe:
            if stock not in self.positions_data:
                self._handle_entry_for_stock(stock, context)

    def _handle_exit_for_stock(self, stock, context):
        hist = context.data.history(stock, ['open', 'high', 'low', 'close'], self.params['sma_long_window'] + 5, context.dt)
        if hist.empty: return
        current_price = hist['close'].iloc[-1]
        # 确保调用时获取单个值 (默认行为)
        sma_long = calculate_sma(hist['close'], self.params['sma_long_window'])
        rsi = calculate_rsi(hist['close'], self.params['rsi_window'])
        stop_loss_price = self.positions_data[stock]['stop_loss']
        
        exit_reason = None
        if current_price <= stop_loss_price: exit_reason = "Stop-Loss"
        elif current_price < sma_long: exit_reason = "Trend Reversal"
        elif rsi > self.params['rsi_profit_target']: exit_reason = "RSI Take-Profit"
        
        if exit_reason:
            order(stock, 0, style='target')
            log(f"[{context.dt.date()}] EXIT {stock}: {exit_reason} @ {current_price:.2f}")
            del self.positions_data[stock]
    
    def _handle_entry_for_stock(self, stock, context):
        hist = context.data.history(stock, ['open', 'high', 'low', 'close'], self.params['sma_long_window'] + 5, context.dt)
        if len(hist) < self.params['sma_long_window']: return
        current_price = hist['close'].iloc[-1]
        # 确保调用时获取单个值 (默认行为)
        sma_long = calculate_sma(hist['close'], self.params['sma_long_window'])
        rsi = calculate_rsi(hist['close'], self.params['rsi_window'])
        atr = calculate_atr(hist['high'], hist['low'], hist['close'], self.params['atr_window'])

        if any(pd.isna(v) for v in [current_price, sma_long, rsi, atr]) or atr <= 0: return
        if not (current_price > sma_long and rsi < self.params['rsi_pullback_level']): return

        risk_per_share = self.params['atr_stop_multiplier'] * atr
        portfolio_value = context.portfolio.update_total_value(context.data.get_current_prices(context.dt))
        position_risk_value = portfolio_value * self.params['risk_per_trade_percent']
        
        shares_to_buy = int(position_risk_value / risk_per_share)
        shares_to_buy = (shares_to_buy // 100) * 100

        if shares_to_buy == 0: return
        if context.portfolio.cash < shares_to_buy * current_price:
            shares_to_buy = int(context.portfolio.cash / current_price // 100) * 100
            if shares_to_buy == 0: return

        order(stock, shares_to_buy)
        stop_loss_price = current_price - risk_per_share
        self.positions_data[stock] = {'stop_loss': stop_loss_price}
        log(f"[{context.dt.date()}] ENTRY {stock}: Buy {shares_to_buy} @ {current_price:.2f}, SL: {stop_loss_price:.2f}")

if __name__ == '__main__':
    cerebro = Cerebro()
    cerebro.add_strategy(SystematicRiskParityStrategy)
    cerebro.set_universe(get_hs300_components_mock())
    cerebro.set_commission(commission=0.00025, stamp_duty=0.0005, min_commission=5.0)

    cerebro.configure(
        start_date='2019-01-01',
        end_date='2024-12-31',
        initial_cash=1000000,
        data_source='real',
        data_path='./data'
    )

    # 运行回测并获取增强的分析结果对象
    results = cerebro.run()

    # 使用新的报告和绘图功能
    if results:
        # 1. 生成详细的文字报告并保存
        results.generate_report()
        
        # 2. 绘制资金曲线和回撤图 (不立即显示，仅保存)
        results.plot_equity_curve(show=False, save=True)

        # 3. 绘制所有交易过的股票的K线图并保存
        results.plot_trades_on_candlestick(save=True)

        print("\n[--- All reports and charts have been generated successfully! ---]")